# # SPDX-License-Identifier: Apache-2.0
import os 
import torch
import multiprocessing as mp

from vllm import LLM, SamplingParams
from vllm import SamplingParams


os.environ["VLLM_USE_V1"]  = "0"


prompt = "很高兴见到你, " 
model  = "/data/yangxianpku/models/Qwen/Qwen2.5-7B-Instruct"

def main():
    # 初始化LLM引擎
    llm = LLM(model=model,
            max_model_len=4096,
            tensor_parallel_size=2,
            trust_remote_code=True,
            enable_chunked_prefill=False,
            dtype=torch.bfloat16,
            enforce_eager=True,
            gpu_memory_utilization=0.8
          )
    
    # 进行推理
    sampling_params = SamplingParams(
                                temperature=0, 
                                top_p=0.95, 
                                max_tokens=100
                              )
    outputs = llm.generate(prompt, sampling_params)
    
    for output in outputs:
        print(f"Generated text: {output.outputs[0].text}")

if __name__ == "__main__":
    # 设置多进程启动方法
    mp.set_start_method('spawn', force=True)
    main()